import gmsh
import sys

class PlasamHMesh:
    def __init__(self, efit):
        self.efit = efit

    def generate_mesh(self, mesh_length = 0.05):
        
        gmsh.initialize()
        gmsh.model.add("gmsh_occ_mesh")

        self.mesh_length = mesh_length

        #define first wall, including upper divertor
        self.first_wall_point_list = []
        for i in range(self.efit.finalwall_fw_outer.shape[0]):
            point = gmsh.model.occ.addPoint(self.efit.finalwall_fw_outer[i, 0], self.efit.finalwall_fw_outer[i, 1], 0, self.mesh_length)
            self.first_wall_point_list.append(point)
            print("point index = ", point)

        for i in range(self.efit.finalwall_div_ur.shape[0]):
            point = gmsh.model.occ.addPoint(self.efit.finalwall_div_ur[i, 0], self.efit.finalwall_div_ur[i, 1], 0, self.mesh_length)
            self.first_wall_point_list.append(point)
            print("point index = ", point)

        for i in range(self.efit.finalwall_dome_u.shape[0]):
            point = gmsh.model.occ.addPoint(self.efit.finalwall_dome_u[i, 0], self.efit.finalwall_dome_u[i, 1], 0, self.mesh_length)
            self.first_wall_point_list.append(point)
            print("point index = ", point)

        for i in range(self.efit.finalwall_div_ul.shape[0]):
            point = gmsh.model.occ.addPoint(self.efit.finalwall_div_ul[i, 0], self.efit.finalwall_div_ul[i, 1], 0, self.mesh_length)
            self.first_wall_point_list.append(point)
            print("point index = ", point)

        for i in range(self.efit.finalwall_fw_inner.shape[0]):
            point = gmsh.model.occ.addPoint(self.efit.finalwall_fw_inner[i, 0], self.efit.finalwall_fw_inner[i, 1], 0, self.mesh_length)
            self.first_wall_point_list.append(point)
            print("point index = ", point)


        self.first_wall_line_list = []
        for i in range(len(self.first_wall_point_list)-1):
            line = gmsh.model.occ.addLine(self.first_wall_point_list[i], self.first_wall_point_list[i+1])
            self.first_wall_line_list.append(line)


        #define divertor_ll
        self.divertor_ll_point_list = []
        for i in range(self.efit.finalwall_div_ll.shape[0]):
            point = gmsh.model.occ.addPoint(self.efit.finalwall_div_ll[i, 0], self.efit.finalwall_div_ll[i, 1], 0, self.mesh_length)
            self.divertor_ll_point_list.append(point)

        print("self.first_wall_point_list  len: ", len(self.first_wall_point_list))
        print("self.divertor_ll_point_list len: ", len(self.divertor_ll_point_list))
        print("self.efit.finalwall_div_ll ", self.efit.finalwall_div_ll)
        self.divertor_ll_line_list = []
        line = gmsh.model.occ.addLine(self.first_wall_point_list[-1], self.divertor_ll_point_list[0])
        self.divertor_ll_line_list.append(line)
        for i in range(len(self.divertor_ll_point_list)-1):
            line = gmsh.model.occ.addLine(self.divertor_ll_point_list[i], self.divertor_ll_point_list[i+1])
            self.divertor_ll_line_list.append(line)

        #define dome_l
        self.dome_l_point_list = []
        for i in range(self.efit.finalwall_dome_l.shape[0]):
            point = gmsh.model.occ.addPoint(self.efit.finalwall_dome_l[i, 0], self.efit.finalwall_dome_l[i, 1], 0, self.mesh_length)
            self.dome_l_point_list.append(point)


        self.dome_l_line_list = []
        line = gmsh.model.occ.addLine(self.divertor_ll_point_list[-1], self.dome_l_point_list[0])
        self.dome_l_line_list.append(line)
        for i in range(len(self.dome_l_point_list)-1):
            line = gmsh.model.occ.addLine(self.dome_l_point_list[i], self.dome_l_point_list[i+1])
            self.dome_l_line_list.append(line)


        #define divertor_lr
        self.divertor_lr_point_list = []
        for i in range(self.efit.finalwall_div_lr.shape[0]):
            point = gmsh.model.occ.addPoint(self.efit.finalwall_div_lr[i, 0], self.efit.finalwall_div_lr[i, 1], 0, self.mesh_length)
            self.divertor_lr_point_list.append(point)


        self.divertor_lr_line_list = []
        line = gmsh.model.occ.addLine(self.dome_l_point_list[-1], self.divertor_lr_point_list[0])
        self.divertor_lr_line_list.append(line)
        for i in range(len(self.divertor_lr_point_list)-1):
            line = gmsh.model.occ.addLine(self.divertor_lr_point_list[i], self.divertor_lr_point_list[i+1])
            self.divertor_lr_line_list.append(line)
        line = gmsh.model.occ.addLine(self.divertor_lr_point_list[-1], self.first_wall_point_list[0])
        self.divertor_lr_line_list.append(line)


        self.first_wall_curve_loop = gmsh.model.occ.addCurveLoop(self.first_wall_line_list + self.divertor_ll_line_list + self.dome_l_line_list + self.divertor_lr_line_list)

        self.first_wall_plane = gmsh.model.occ.addPlaneSurface([self.first_wall_curve_loop])

        
        #define inner edge
        self.inner_edge_point_list = []
        for i in range(len(self.efit.inner_edge)):
            if i > 0 and self.efit.inner_edge[i, 0] == self.efit.inner_edge[i-1, 0] and self.efit.inner_edge[i, 1] == self.efit.inner_edge[i-1, 1]:
                pass
            elif i == len(self.efit.inner_edge) - 1 and self.efit.inner_edge[i, 0] == self.efit.inner_edge[0, 0] and self.efit.inner_edge[i, 1] == self.efit.inner_edge[0, 1]:
                pass
            else:
                point = gmsh.model.occ.addPoint(self.efit.inner_edge[i, 0], self.efit.inner_edge[i, 1], 0, self.mesh_length)
                self.inner_edge_point_list.append(point)


        self.inner_edge_line_list = []
        for i in range(len(self.inner_edge_point_list)):
            #print("line list ", self.efit.inner_edge[i, 0], self.efit.inner_edge[i, 1], self.efit.inner_edge[i+1, 0], self.efit.inner_edge[i+1, 1])
            if i == len(self.inner_edge_point_list) - 1:
                line = gmsh.model.occ.addLine(self.inner_edge_point_list[-1], self.inner_edge_point_list[0])
            else:
                line = gmsh.model.occ.addLine(self.inner_edge_point_list[i], self.inner_edge_point_list[i+1])
            self.inner_edge_line_list.append(line)
            #print("line list ", i, line)


        self.inner_edge_curve_loop = gmsh.model.occ.addCurveLoop(self.inner_edge_line_list)

        self.inner_edge_plane = gmsh.model.occ.addPlaneSurface([self.inner_edge_curve_loop])
        


        #cut
        #self.sim_plane = gmsh.model.occ.remove([(2, self.inner_edge_plane)])
        #self.sim_plane = gmsh.model.occ.cut([(2, self.first_wall_plane)], [(2, self.inner_edge_plane)])
        self.sim_plane = gmsh.model.occ.fragment([(2, self.first_wall_plane)], [(2, self.inner_edge_plane)])
        gmsh.model.occ.remove([(2, self.inner_edge_plane)], recursive = True)
        #print("self.sim_plane ", self.sim_plane)
        

        gmsh.model.occ.synchronize()

        #print("self.first_wall_line_list ", self.first_wall_line_list)

        #define physical group
        self.first_wall_physical_group = gmsh.model.addPhysicalGroup(1, self.first_wall_line_list)
        gmsh.model.setPhysicalName(1, self.first_wall_physical_group, "first_wall")

        self.divertor_ll_physical_group = gmsh.model.addPhysicalGroup(1, self.divertor_ll_line_list)
        gmsh.model.setPhysicalName(1, self.divertor_ll_physical_group, "divertor_ll")

        self.dome_l_physical_group = gmsh.model.addPhysicalGroup(1, self.dome_l_line_list)
        gmsh.model.setPhysicalName(1, self.dome_l_physical_group, "dome_l")

        self.divertor_lr_physical_group = gmsh.model.addPhysicalGroup(1, self.divertor_lr_line_list)
        gmsh.model.setPhysicalName(1, self.divertor_lr_physical_group, "divertor_lr")

        self.inner_edge_physical_group = gmsh.model.addPhysicalGroup(1, self.inner_edge_line_list)
        gmsh.model.setPhysicalName(1, self.inner_edge_physical_group, "inner_edge")

        #self.sim_plane_physical_group = gmsh.model.addPhysicalGroup(2, [self.first_wall_plane])
        #self.sim_plane_physical_group = gmsh.model.addPhysicalGroup(2, [self.sim_plane])
        self.sim_plane_physical_group = gmsh.model.addPhysicalGroup(2, [3])
        gmsh.model.setPhysicalName(2, self.sim_plane_physical_group, "sim_plane")
        





        gmsh.model.occ.synchronize()
        gmsh.model.mesh.generate(2)




        #write mesh
        #gmsh.option.setNumber("Mesh.SaveAll", 1)
        gmsh.option.setNumber("Mesh.MshFileVersion", 2.2)
        #gmsh.write("hl2a.msh")

        # Launch the GUI to see the results:
        if '-nopopup' not in sys.argv:
            gmsh.fltk.run()

    def save_mesh(self, mesh_name = "hl2a.msh"):
        gmsh.write("model.brep")
        gmsh.write(mesh_name)
        gmsh.finalize()
